package core

import (
	"encoding/binary"
	"fmt"
)

// IPHeader IPv4头部结构
type IPHeader struct {
	Version                uint8
	HeaderLength           uint8
	DifferentiatedServices uint8
	TotalLength            uint16
	Identification         uint16
	Flags                  uint8
	FragmentOffset         uint16
	TimeToLive             uint8
	Protocol               uint8
	Checksum               uint16
	SourceAddress          uint32
	DestinationAddress     uint32
	Options                []byte
}

// IPPackage 完整的IP数据包
type IPPackage struct {
	Header  IPHeader
	Payload []byte
}

// ParseIPPackage 解析IP数据包
func ParseIPPackage(data []byte) (*IPPackage, error) {
	if len(data) < 20 {
		return nil, fmt.Errorf("数据包太短，无法解析IP头部")
	}

	pkg := &IPPackage{}

	// 解析版本和头长度（第1个字节）
	versionAndIHL := data[0]
	pkg.Header.Version = (versionAndIHL >> 4) & 0x0F
	pkg.Header.HeaderLength = versionAndIHL & 0x0F

	// 检查版本是否为IPv4
	if pkg.Header.Version != 4 {
		return nil, fmt.Errorf("不是IPv4数据包，版本: %d", pkg.Header.Version)
	}

	// 计算实际头部长度
	headerLengthBytes := pkg.Header.HeaderLength * 4
	if len(data) < int(headerLengthBytes) {
		return nil, fmt.Errorf("数据包长度不足，期望至少%d字节", headerLengthBytes)
	}

	// 服务类型（第2个字节）
	pkg.Header.DifferentiatedServices = data[1]

	// 总长度（第3-4字节）
	pkg.Header.TotalLength = binary.BigEndian.Uint16(data[2:4])

	// 标识（第5-6字节）
	pkg.Header.Identification = binary.BigEndian.Uint16(data[4:6])

	// 标志和片段偏移（第7-8字节）
	flagsAndOffset := binary.BigEndian.Uint16(data[6:8])
	pkg.Header.Flags = uint8((flagsAndOffset >> 13) & 0x07)
	pkg.Header.FragmentOffset = flagsAndOffset & 0x1FFF

	// 生存时间（第9字节）
	pkg.Header.TimeToLive = data[8]

	// 协议（第10字节）
	pkg.Header.Protocol = data[9]

	// 头校验和（第11-12字节）
	pkg.Header.Checksum = binary.BigEndian.Uint16(data[10:12])

	// 源IP地址（第13-16字节）
	pkg.Header.SourceAddress = binary.BigEndian.Uint32(data[12:16])

	// 目标IP地址（第17-20字节）
	pkg.Header.DestinationAddress = binary.BigEndian.Uint32(data[16:20])

	// 选项（如果有）
	if headerLengthBytes > 20 {
		pkg.Header.Options = make([]byte, headerLengthBytes-20)
		copy(pkg.Header.Options, data[20:headerLengthBytes])
	}

	// 计算payload长度：PayloadLength = Header.TotalLength - 4 * Header.HeaderLength
	payloadLength := int(pkg.Header.TotalLength) - int(pkg.Header.HeaderLength)*4
	if payloadLength > 0 && len(data) >= int(headerLengthBytes)+payloadLength {
		pkg.Payload = make([]byte, payloadLength)
		copy(pkg.Payload, data[headerLengthBytes:headerLengthBytes+uint8(payloadLength)])
	}

	return pkg, nil
}

// GetPayloadLength 获取payload长度
func (p *IPPackage) GetPayloadLength() int {
	return int(p.Header.TotalLength) - int(p.Header.HeaderLength)*4
}
